import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

plt.style.use('seaborn-white grid')


# 循环访问行，对于每一行，我们查询数据记录以查看有多少行与准标识符的值匹配。
# 如果任何组中的行数小于k，则数据帧不满足该值k-Anonymity，我们返回 False。
# 请注意，在这个简单的定义中，我们认为所有列都包含准标识符;要将检查限制为所有列的子集，我们需要将表达式替换为其他内容。
# 当k=1的时候是满足条件的返回true
def isKAnonymized(df, k):
    for index, row in df.iterrows():
        query = ' & '.join([f'{col} == {row[col]}' for col in df.columns])
        rows = df.query(query)
        if rows.shape[0] < k:  # len(rows)
            return False
    return True


# 这个函数里面的y表示表中实际数据，比如42。depths[x.name]表示当前属性下设定的映射值，下面在depths中将各个属性设为1，也就是对42/10再*10得到40。整个函数就是对所有数字向下取到10的倍数。
def generalize(df, depths):
    return df.apply(lambda x: x.apply(lambda y: int(int(y / (10 ** depths[x.name])) * (10 ** depths[x.name]))))


# 新的数据集
adult_data = pd.read_csv("PUMAS.csv")
print(adult_data.head())

# 当前情况k = 2的 k-Anonymity不满足，k = 1的k-Anonymity已经满足。
df = adult_data[['age', 'educ']]
df.columns = ['ages', 'educs']
print(isKAnonymized(df.head(100), 1))
# isKAnonymized(df.head(100), 2)

# 老方法，调用generalize对数据进行处理。结果仍然不满足
depths = {
    'ages': 1,
    'educs': 1
}
df2 = generalize(df.head(1000), depths)
print(isKAnonymized(df2, 2))

# 删除异常值（Removing Outliers），使用Numpy的方法来执行此裁剪。我们将年龄裁剪为50岁或以下，而将教育水平单独保留（通过将其裁剪为非常大的值）。
# DataFrame.clip(lower=None, upper=None, axis=None, inplace=False, *args, **kwargs)
depths = {
    'ages': 1,
    'educs': 1
}
dfp = df.clip(upper=np.array([50, 10000000000000]), axis='columns')
df2 = generalize(dfp.head(500), depths)
print(isKAnonymized(df2, 4))  # 满足k = 4 的k-Anonymity
'''
   age  sex  educ  race   income  married  pid
0   59    1     9     1      0.0        1    1
1   31    0     1     3  17000.0        0    2
2   36    1    11     1      0.0        1    3
3   54    1    11     1   9100.0        1    4
4   39    0     5     3  37000.0        0    5
True
False
True
'''